package cn.gjsm.utils;

import cn.gjsm.pojos.OpenAi;
import com.theokanning.openai.service.OpenAiService;
import com.theokanning.openai.completion.CompletionChoice;
import com.theokanning.openai.completion.CompletionRequest;
import org.jetbrains.annotations.NotNull;
import org.springframework.util.StringUtils;

import java.time.Duration;
import java.util.Arrays;
import java.util.List;

public class OpenAiClient {
    public static final int MAX_TOKENS = 1000;
    public static String OPENAPI_TOKEN = "";
    public static Integer TIMEOUT = null;

    /**
     * 获取ai
     *
     * @param openAi
     * @param prompt
     * @return
     */
    public static List<CompletionChoice> getAiResult(OpenAiService service, OpenAi openAi, String prompt) {
        //OpenAiService service = getOpenAiService();
        CompletionRequest.CompletionRequestBuilder builder = CompletionRequest.builder()
                .model(openAi.getModel())
                .prompt(prompt)
                .temperature(openAi.getTemperature())
                .maxTokens(MAX_TOKENS)
                .topP(openAi.getTopP())
                .frequencyPenalty(openAi.getFrequencyPenalty())
                .presencePenalty(openAi.getPresencePenalty());
        if (!StringUtils.isEmpty(openAi.getStop())) {
            builder.stop(Arrays.asList(openAi.getStop().split(",")));
        }
        var request = builder.build();
        var result= service.createCompletion(request);

        return result.getChoices();
    }

    @NotNull
    public static OpenAiService getOpenAiService() {
        if (TIMEOUT == null || TIMEOUT < 100) {
            TIMEOUT = 100;
        }

        var duration = Duration.ofSeconds(TIMEOUT);
        OpenAiService service = new OpenAiService(OPENAPI_TOKEN, duration);
        return service;
    }
}
